# Copyright (c) 2020 cgsdfc
#
# This software is released under the MIT License.
# https://opensource.org/licenses/MIT
"""
Generate things defined in cppdefs.py for the formmating code.

The generating code is deliberately separated from what is generated
so that the asdl module is isolated from the other part.
"""

import itertools
from typing import Union, Sequence, Callable
from astlib.asdl import (VisitorBase, Constructor, Module, Field, Product, Sum,
                         Type, builtin_types, parse as ParseAsdl)

from astlib.cc_types import (ConstructorNode, TypeNode, ProductNode, SumEnum,
                             KindEnum, Member, Builtin, Scalar, Object,
                             AsdlSeq, Location, NumArgs, Hierarchy, Config)


def GetConstructorNames(sum: Sum):
    return [cons.name for cons in sum.types]


def IsScalar(asdl_name):
    return 'int' == asdl_name


def CreateBuiltin(asdl_name):
    assert IsBuiltinType(asdl_name)
    return (Builtin.CreateScalar
            if IsScalar(asdl_name) else Builtin.CreateObject)(asdl_name)


def IsBuiltinType(asdl_name):
    return asdl_name in builtin_types


def IsSimpleSum(node: Sum):
    return all(not t.fields for t in node.types)


def HasAttributes(sum_or_product: Union[Sum, Product]):
    return len(sum_or_product.attributes) > 0


def GetLocationInfo(attributes):
    field_names = frozenset(f.name for f in attributes)
    has_location = all(n in field_names for n in Location.LOCATION_FIELDS)
    has_end_location = all(n in field_names
                           for n in Location.LOCATION_END_FIELDS)
    return dict(has_location=has_location, has_end_location=has_end_location)


def IsSequence(field: Field):
    return field.seq


def IsOptional(field: Field):
    return field.opt


def GetNumArgs(field: Field):
    if IsOptional(field):
        return NumArgs.OPTIONAL
    if IsSequence(field):
        return NumArgs.SEQUENCE
    return NumArgs.REQUIRED


class Generator(VisitorBase):
    def __init__(self):
        super().__init__()
        self._types = {}
        self._type_nodes = {}
        self._sequence_types = {}
        self._builtin_types = {}
        self._deferred_tasks = []

    def FindSeqType(self, element_type: str):
        if element_type in self._sequence_types:
            return self._sequence_types[element_type]
        asdl_name = element_type
        element_type = self.FindNonSeqType(asdl_name)
        seq = AsdlSeq(asdl_name, element_type)
        return self._sequence_types.setdefault(asdl_name, seq)

    def FindBuiltinType(self, asdl_name):
        if asdl_name in self._builtin_types:
            return self._builtin_types[asdl_name]
        return self._builtin_types.setdefault(asdl_name,
                                              CreateBuiltin(asdl_name))

    def FindNonSeqType(self, asdl_name):
        if IsBuiltinType(asdl_name):
            return self.FindBuiltinType(asdl_name)
        return self._types[asdl_name]

    def FindTypeByField(self, field: Field):
        asdl_name = field.type
        if IsSequence(field):
            return self.FindSeqType(asdl_name)
        return self.FindNonSeqType(asdl_name)

    def AddTypeNode(self, asdl_name, type_node: TypeNode):
        assert asdl_name not in self._type_nodes
        self._type_nodes[asdl_name] = type_node

    def AddType(self, asdl_name, type):
        table = self._type_nodes if isinstance(type, KindEnum) else self._types
        assert asdl_name not in table
        table[asdl_name] = type

    def AddDeferredTask(self, task: Callable[[], None]):
        self._deferred_tasks.append(task)

    def visitModule(self, mod: Module):
        for dfn in mod.dfns:
            self.visit(dfn)

        for task in self._deferred_tasks:
            task()

    def visitType(self, type: Type):
        self.visit(type.value, str(type.name))

    def visitSum(self, sum: Sum, name: str):
        cons_names = GetConstructorNames(sum)

        if IsSimpleSum(sum):
            sum_enum = SumEnum(name, cons_names)
            self.AddType(name, sum_enum)
        else:
            kind_enum = KindEnum(name, cons_names)
            self.AddType(name, kind_enum)
            type_node = TypeNode(name,
                                 kind_enum=kind_enum,
                                 **GetLocationInfo(sum.attributes))
            self.AddType(name, type_node)

            for cons in sum.types:
                self.visit(cons, type_node)

    def CreateMembers(self, fields: Sequence[Field]):
        return tuple(self.visit(f) for f in fields)

    def visitConstructor(self, cons: Constructor, parent_node: TypeNode):
        asdl_name = cons.name
        cons_node = ConstructorNode(asdl_name,
                                    parent=parent_node,
                                    members=None)
        self.AddType(asdl_name, type=cons_node)

        def Task():
            cons_node.members = self.CreateMembers(cons.fields)

        self.AddDeferredTask(Task)

    def visitProduct(self, prod: Product, asdl_name):
        product_node = ProductNode(asdl_name,
                                   members=None,
                                   **GetLocationInfo(prod.attributes))
        self.AddType(asdl_name, product_node)

        def Task():
            product_node.members = self.CreateMembers(prod.fields)

        self.AddDeferredTask(Task)

    def visitField(self, field: Field) -> Member:
        type = self.FindTypeByField(field)
        asdl_name = field.name
        nargs = GetNumArgs(field)
        return Member(type, asdl_name, nargs)

    def MergeTables(self) -> list:
        tables = (self._types, self._builtin_types, self._sequence_types,
                  self._type_nodes)
        type_list = []
        for tab in tables:
            type_list.extend(tab.values())
        return type_list

    def GenerateHierarchy(self, mod: Module):
        self.visit(mod)
        type_list = self.MergeTables()
        return Hierarchy(type_list=type_list)


def CreateTemplateData(filename) -> dict:
    mod = ParseAsdl(filename)
    gen = Generator()
    hierarchy = gen.GenerateHierarchy(mod)
    config = Config.Create(hierarchy)
    return dict(hierarchy=hierarchy, config=config)


if __name__ == "__main__":
    from argparse import ArgumentParser
    import pickle

    p = ArgumentParser()
    p.add_argument('-i', help='input Python.asdl')
    p.add_argument('-o', help='output data.pkl')
    args = p.parse_args()

    data = CreateTemplateData(str(args.i))
    with open(str(args.o), 'wb') as f:
        pickle.dump(data, f)
